# coding = utf-8

'''
分析肿瘤灰度值
'''

import os
import numpy as np
import pydicom
import matplotlib.pyplot as plt

def get_pixels_hu(scans):
    image = np.stack([s.pixel_array for s in scans])
    # Convert to int16 (from sometimes int16),
    # should be possible as values should always be low enough (<32k)
    image = image.astype(np.int16)

    # Set outside-of-scan pixels to 1
    # The intercept is usually -1024, so air is approximately 0
    image[image == -2000] = 0

    # Convert to Hounsfield units (HU)
    intercept = scans[0].RescaleIntercept
    slope = scans[0].RescaleSlope

    if slope != 1:
        image = slope * image.astype(np.float64)
        image = image.astype(np.int16)

    image += np.int16(intercept)
    return np.array(image, dtype=np.int16)

def readImage(path):
    index_2_image = {}
    for item in os.listdir(path):
        file_name = os.path.join(path, item)
        dcm = pydicom.dcmread(file_name)
        hu = get_pixels_hu([dcm])
        hu = hu.squeeze()
        index_2_image[item.split("_")[1].zfill(3)] = hu
    return index_2_image

def readMask(path):
    index_2_mask = {}
    liver_mask_path = os.path.join(path, "liver")
    tumor_mask_list = []
    for name in os.listdir(path):
        if "livertumor" in name:
            tumor_mask_list.append(os.path.join(path,name))
    for file_index in os.listdir(liver_mask_path):
        liver_mask_file = os.path.join(liver_mask_path, file_index)
        liver_mask = pydicom.dcmread(liver_mask_file)
        mask = liver_mask.pixel_array
        mask_division = np.max(mask)
        if mask_division == 0:
            mask_division = 255
        mask = mask/mask_division
        mask = mask.astype(np.uint8)
        for tumor_mask_path in tumor_mask_list:
            tumor_mask_file = os.path.join(tumor_mask_path, file_index)
            assert os.path.exists(tumor_mask_file), tumor_mask_file+" not exist"
            tumor_mask = pydicom.dcmread(tumor_mask_file)
            tumor_mask = tumor_mask.pixel_array
            tumor_division = np.max(tumor_mask)
            if tumor_division == 0:
                tumor_division = 255
            tumor_mask = tumor_mask / tumor_division
            #print(tumor_division)
            tumor_mask = tumor_mask.astype(np.uint8)
            tumor_size = (tumor_mask == 1).sum()
            liver_size = (mask == 1).sum()
            inter = (mask[tumor_mask == 1] == 1).sum()
            #assert (tumor_size-inter) == 0, "{},liver size:{}, tumor size:{}, inter:{}, diff:{}".format(tumor_mask_file,liver_size, tumor_size, inter,
            #                                                                                         tumor_size-inter)
            mask[tumor_mask == 1] = 2
        index_2_mask[file_index.split("_")[1].zfill(3)] = mask
    return index_2_mask

def calcute():
    root = "/datasets/3Dircadb/3Dircadb1"
    chenkung_root = "/datasets/3Dircadb/chengkung"
    for patient_id in os.listdir(root):
        patient_path = os.path.join(root, patient_id)
        image_path = os.path.join(patient_path, "PATIENT_DICOM")
        mask_path = os.path.join(patient_path, "MASKS_DICOM")
        case_id = "case_{}".format(str(int(patient_id.split(".")[1]) - 1).zfill(5))
        index_2_image = readImage(image_path)
        index_2_mask = readMask(mask_path)

        his_data = []
        tumor_data = []
        for key in sorted(index_2_image.keys()):
            image = index_2_image[key]
            mask = index_2_mask[key]
            if np.max(mask) == 0:
                continue
            for row in range(mask.shape[0]):
                for column in range(mask.shape[1]):
                    if mask[row][column] == 0:
                        continue
                    if image[row][column] < -250 or image[row][column] > 250:
                        continue
                    if mask[row][column] == 1:
                        his_data.append(image[row][column])
                    if mask[row][column] == 2:
                        tumor_data.append(image[row][column])
        liver = np.array(his_data)
        tumor = np.array(tumor_data)

        fig, ax1 = plt.subplots()
        ax2 = ax1.twinx()
        ax1.hist([liver, tumor], bins=100, color=['b','r'])
        n, bins, patches = ax1.hist([liver, tumor], bins=100)
        ax1.cla()  # clear the axis

        width = (bins[1] - bins[0]) * 0.4
        bins_shifted = bins + width
        ax1.bar(bins[:-1], n[0], width, align='edge', color='b')
        ax2.bar(bins_shifted[:-1], n[1], width, align='edge', color='r')

        colors = ['b', 'r']

        ax1.set_ylabel("liver hu Count", color=colors[0])
        ax2.set_ylabel("tumor hu Count", color=colors[1])
        ax1.tick_params('y', colors=colors[0])
        ax2.tick_params('y', colors=colors[1])
        plt.tight_layout()
        plt.title("total")





        plt.show()








if __name__ == '__main__':
    calcute()